/**********************************************************************
 * SailmiPIM - modular cross-platform feature Personal Information Manager Tools.
 * Copyright(C) 2006-2014  Georg Rudoy
 *
 * Distributed under the Boost Software License, Version 1.0.
 *(See accompanying file LICENSE or copy at https://www.boost.org/LICENSE_1_0.txt)
 **********************************************************************/

#pragma once

#include <coroutine>
#include <utility>
#include "finalsuspender.h"
#include "taskfwd.h"

namespace SailmiPIM::Util
{
	namespace detail
	{
		template<typename R>
		struct PromiseRet
		{
			using ReturnType_t = R;

			constexpr static bool IsVoid = false;

			std::optional<R> Ret_;

			template<typename U>
			void return_value(U&& val)
			{
				Ret_.emplace(std::forward<U>(val));
			}
		};

		template<>
		struct PromiseRet<void>
		{
			constexpr static bool IsVoid = true;

			bool Done_ = false;

			void return_void() noexcept
			{
				Done_ = true;
			}
		};

		struct EitherFailureAbort : std::exception {};

		template<typename Promise>
		struct TaskAwaiter
		{
			using Handle_t = std::coroutine_handle<Promise>;
			Handle_t Handle_;

			bool await_ready() const noexcept
			{
				const auto& promise = Handle_.promise();
				if(promise.Exception_)
					return true;

				if constexpr(Promise::IsVoid)
					return promise.Done_;
				else
					return static_cast<bool>(promise.Ret_);
			}

			void await_suspend(std::coroutine_handle<> handle)
			{
				Handle_.promise().WaitingHandles_.push_back(handle);
			}

			auto await_resume() const
			{
				const auto& promise = Handle_.promise();
				if(promise.Exception_)
					try
					{
						std::rethrow_exception(promise.Exception_);
					}
					catch(EitherFailureAbort)
					{
					}

				if constexpr(!Promise::IsVoid)
					return *promise.Ret_;
			}
		};
	}

	template<typename R, template<typename> typename... Extensions>
	class Task
	{
	public:
		struct promise_type;
	private:
		using Handle_t = std::coroutine_handle<promise_type>;
		Handle_t Handle_;
	public:
		struct promise_type : detail::PromiseRet<R>
							, Extensions<promise_type>...
		{
			size_t Refs_ = 1; // TODO make thread-safe
			QVector<std::coroutine_handle<>> WaitingHandles_ {};
			std::exception_ptr Exception_ {};

			auto GetAddress() { return Handle_t::from_promise(*this).address(); }

			Task get_return_object()
			{
				return Task { Handle_t::from_promise(*this) };
			}

			std::suspend_never initial_suspend() const noexcept { return {}; }

			auto final_suspend() noexcept
			{
				([this]
				{
					using Base = Extensions<promise_type>;
					if constexpr(requires(Base t) { t.FinalSuspend(); })
						Base::FinalSuspend();
				}(), ...);
				return detail::FinalSuspender<promise_type> { *this };
			}

			void unhandled_exception()
			{
				Exception_ = std::current_exception();
			}

			void IncRef()
			{
				++Refs_;
			}

			void DecRef()
			{
				if(!--Refs_)
					Handle_t::from_promise(*this).destroy();
			}
		};

		using ResultType_t = R;

		template<typename RR>
		using ReplaceResult_t = Task<RR, Extensions...>;

		explicit Task(const std::coroutine_handle<promise_type>& handle)
		: Handle_ { handle }
		{
			if(handle)
				handle.promise().IncRef();
		}

		~Task() noexcept
		{
			if(Handle_)
				Handle_.promise().DecRef();
		}

		Task(const Task& other)
		: Handle_ { other.Handle_ }
		{
			if(Handle_)
				Handle_.promise().IncRef();
		}

		Task& operator=(const Task& other)
		{
			Task task { other };
			*this = std::move(task);
			return *this;
		}

		Task(Task&& other) noexcept
		{
			std::swap(Handle_, other.Handle_);
		}

		Task& operator=(Task&& other) noexcept
		{
			std::swap(Handle_, other.Handle_);
			return *this;
		}

		auto operator co_await() const noexcept
		{
			return detail::TaskAwaiter<promise_type> { Handle_ };
		}
	};

	namespace detail
	{
		template<typename R, template<typename> typename... Extensions>
		struct GetPromise
		{
			using Promise = typename Task<R, Extensions...>::promise_type;
			Promise *Promise_ = nullptr;

			bool await_ready() const noexcept { return false; }

			bool await_suspend(std::coroutine_handle<Promise> handle) const noexcept
			{
				Promise_ = &handle.promise();
				return false;
			}

			decltype(auto) await_resume() const noexcept
			{
				return *Promise_;
			}
		};
	}
}
